import {Annotation, END, START, StateGraph} from "@langchain/langgraph";
import {saveArrayBufferAsImage} from "../tool.mjs";


const StateAnnotation = Annotation.Root({
    aggregate: Annotation({
        reducer: (x, y) => {
            console.log(x, y)
           return  x.concat(y)
        }
    }),
    which: Annotation({
        reducer: (x, y) => {
            console.log(x, y)
            return (y ?? x);
        },
    })
})


const nodeA2 = state => {
    console.log(`Adding I'm A to ${state.aggregate}`);
    return {
        aggregate: [`I'm A`]
    }
}


const nodeB2 = state => {
    console.log(`Adding I'm B to ${state.aggregate}`);
    return {
        aggregate: [`I'm B`]
    }
}

const nodeC2 = state => {
    console.log(`Adding I'm C to ${state.aggregate}`);
    return {
        aggregate: [`I'm C`]
    }
}


const nodeD2 = state => {
    console.log(`Adding I'm D to ${state.aggregate}`);
    return {
        aggregate: [`I'm D`]
    }
}
const nodeE2 = state => {
    console.log(`Adding I'm E to ${state.aggregate}`);
    return {
        aggregate: [`I'm E`]
    }
}

function routeCDorBc(state) {
    if(state.which === 'cd') {
        return ["c", "d"]
    }
    return ["b", "c"]
}

const builder = new StateGraph(StateAnnotation)
    .addNode("a", nodeA2)
    .addEdge(START, "a")
    .addNode("b", nodeB2)
    .addNode("c", nodeC2)
    .addNode("d", nodeD2)
    .addNode("e", nodeE2)
    .addConditionalEdges("a", routeCDorBc, ['b', 'c', 'd'])
    .addEdge("b", "e")
    .addEdge("c", "e")
    .addEdge("d", "e")
    .addEdge("e", END)


const graph = builder.compile();


const representation =await graph.getGraphAsync()

const image = await representation.drawMermaidPng();
const arrayBuffer = await image.arrayBuffer();

saveArrayBufferAsImage(arrayBuffer, "condition.png")


const baseResult = await graph.invoke({aggregate: [], which: "bc"})

console.log('Base Results:', baseResult)
